#include "lw_mqtt.h"

#ifdef ENABLE_MQTT

#include <cstring>
#include <iostream>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <netdb.h>    // For getaddrinfo
#include <errno.h>    // For errno

#ifdef __cplusplus
extern "C" {
#endif

#include <mqtt.h>
#include <mqtt_pal.h>

#ifdef __cplusplus
}
#endif

LwMqtt::LwMqtt() 
    : initialized_(false), 
      connected_(false),
      port_(1883),
      sockfd_(-1) {
    // Initialize buffers to zero
    memset(sendbuf_, 0, sizeof(sendbuf_));
    memset(recvbuf_, 0, sizeof(recvbuf_));
}

LwMqtt::~LwMqtt() {
    deinit();
}

bool LwMqtt::init(const std::string& host, int port, const std::string& client_id) {
    if (initialized_) {
        std::cerr << "MQTT client already initialized" << std::endl;
        return false;
    }

    host_ = host;
    port_ = port;
    client_id_ = client_id;

    // Resolve hostname to IP address
    struct addrinfo hints, *result;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_INET;       // IPv4
    hints.ai_socktype = SOCK_STREAM; // TCP socket

    std::string port_str = std::to_string(port_);
    int status = getaddrinfo(host_.c_str(), port_str.c_str(), &hints, &result);
    if (status != 0) {
        std::cerr << "Failed to resolve hostname: " << host_ << " (" << gai_strerror(status) << ")" << std::endl;
        return false;
    }

    // Create socket
    sockfd_ = socket(result->ai_family, result->ai_socktype, result->ai_protocol);
    if (sockfd_ < 0) {
        std::cerr << "Failed to create socket: " << strerror(errno) << std::endl;
        freeaddrinfo(result);
        return false;
    }

    // Connect to broker
    if (::connect(sockfd_, result->ai_addr, result->ai_addrlen) < 0) {
        std::cerr << "Failed to connect to broker " << host_ << ":" << port_ 
                  << " (" << strerror(errno) << ")" << std::endl;
        close(sockfd_);
        sockfd_ = -1;
        freeaddrinfo(result);
        return false;
    }

    freeaddrinfo(result);

    // Initialize MQTT client
    mqtt_init(&client_, sockfd_, sendbuf_, sizeof(sendbuf_), recvbuf_, sizeof(recvbuf_), message_callback_wrapper);
    client_.publish_response_callback_state = static_cast<void*>(this); // Fixed: pass this as void pointer

    initialized_ = true;
    return true;
}

void LwMqtt::deinit() {
    if (!initialized_) {
        return;
    }

    disconnect();

    if (sockfd_ >= 0) {
        close(sockfd_);
        sockfd_ = -1;
    }

    initialized_ = false;
    connected_ = false;
}

bool LwMqtt::connect() {
    if (!initialized_) {
        std::cerr << "MQTT client not initialized" << std::endl;
        return false;
    }

    if (connected_) {
        return true;
    }

    std::cout << "Attempting to connect to MQTT broker at " << host_ << ":" << port_ << "..." << std::endl;
    uint8_t connect_flags = MQTT_CONNECT_CLEAN_SESSION;
    enum MQTTErrors err = mqtt_connect(&client_, client_id_.c_str(), NULL, NULL, 0, NULL, NULL, connect_flags, 400);
    if (err != MQTT_OK) {
        std::cerr << "MQTT connection failed with error: " << mqtt_error_str(err) << std::endl;
        return false;
    }

    std::cout << "MQTT connect packet sent, waiting for CONNACK..." << std::endl;
    
    // Wait for CONNACK packet to be received
    mqtt_pal_time_t start_time = MQTT_PAL_TIME();
    mqtt_pal_time_t timeout = 5; // 5 seconds timeout
    
    while (MQTT_PAL_TIME() - start_time < timeout) {
        enum MQTTErrors sync_err = mqtt_sync(&client_);
        
        if (sync_err != MQTT_OK && sync_err != MQTT_ERROR_SEND_BUFFER_IS_FULL) {
            std::cerr << "MQTT sync error: " << mqtt_error_str(sync_err) << std::endl;
        }
        
        // Check if we're really connected by checking for errors
        if (client_.error != MQTT_OK) {
            std::cerr << "MQTT connection failed with error: " << mqtt_error_str(client_.error) << std::endl;
            return false;
        }
        
        // If keep_alive is set, we've successfully connected
        if (client_.keep_alive != 0) {
            std::cout << "MQTT connection established successfully!" << std::endl;
            connected_ = true;
            return true;
        }
        
        usleep(100000); // Wait 100ms before retrying
    }
    
    std::cerr << "MQTT connection timeout - failed to receive CONNACK" << std::endl;
    return false;
}

void LwMqtt::disconnect() {
    if (!connected_) {
        return;
    }

    mqtt_disconnect(&client_);
    connected_ = false;
}

bool LwMqtt::isConnected() const {
    if (!initialized_ || !connected_) {
        return false;
    }

    // Additional check - make sure the client is still connected
    enum MQTTErrors err = mqtt_ping(const_cast<struct mqtt_client*>(&client_));
    if (err != MQTT_OK) {
        std::cerr << "MQTT ping failed: " << mqtt_error_str(err) << ", client error: " << mqtt_error_str(client_.error) << std::endl;
        return false;
    }
    
    return true;
}

bool LwMqtt::publish(const std::string& topic, const std::string& payload, int qos, bool retain) {
    if (!connected_) {
        std::cerr << "MQTT client not connected" << std::endl;
        return false;
    }

    std::cout << "Publishing to topic: " << topic << ", payload: " << payload << ", QoS: " << qos << ", Retain: " << (retain ? "true" : "false") << std::endl;
    
    // Create publish flags from QoS and retain parameters
    uint8_t publish_flags = (qos & 0x03) << 1; // QoS is in bits 1-2
    if (retain) {
        publish_flags |= 0x01; // Retain flag is in bit 0
    }

    enum MQTTErrors err = mqtt_publish(&client_, topic.c_str(), payload.c_str(), payload.length(), publish_flags);
    if (err != MQTT_OK) {
        std::cerr << "MQTT publish failed with error: " << mqtt_error_str(err) << std::endl;
        return false;
    }
    
    std::cout << "MQTT publish packet sent" << std::endl;
    
    // Process outgoing messages immediately
    yield(100);
    
    return true;
}

bool LwMqtt::subscribe(const std::string& topic, int qos) {
    if (!connected_) {
        std::cerr << "MQTT client not connected" << std::endl;
        return false;
    }

    std::cout << "Subscribing to topic: " << topic << " with QoS: " << qos << std::endl;
    enum MQTTErrors err = mqtt_subscribe(&client_, topic.c_str(), qos & 0x03);
    if (err != MQTT_OK) {
        std::cerr << "MQTT subscribe failed with error: " << mqtt_error_str(err) << std::endl;
        return false;
    }
    
    std::cout << "MQTT subscribe packet sent, waiting for SUBACK..." << std::endl;
    
    // Wait for SUBACK packet to be received
    yield(1000);
    
    std::cout << "Subscribe operation completed" << std::endl;
    return true;
}

bool LwMqtt::unsubscribe(const std::string& topic) {
    if (!connected_) {
        std::cerr << "MQTT client not connected" << std::endl;
        return false;
    }

    enum MQTTErrors err = mqtt_unsubscribe(&client_, topic.c_str());
    if (err != MQTT_OK) {
        std::cerr << "MQTT unsubscribe failed with error: " << mqtt_error_str(err) << std::endl;
        return false;
    }

    return true;
}

void LwMqtt::setMessageCallback(const MessageCallback& callback) {
    message_callback_ = callback;
}

void LwMqtt::yield(int timeout_ms) {
    if (!initialized_) {
        return;
    }

    // Process messages for the specified timeout
    mqtt_pal_time_t start_time = MQTT_PAL_TIME();
    while (MQTT_PAL_TIME() - start_time < (timeout_ms/1000.0)) {
        enum MQTTErrors err = mqtt_sync(&client_);
        if (err != MQTT_OK && err != MQTT_ERROR_SEND_BUFFER_IS_FULL) {
            std::cerr << "MQTT sync error: " << mqtt_error_str(err) << std::endl;
        }
        
        // Small delay to prevent excessive CPU usage
        usleep(10000); // 10ms
    }
}

void LwMqtt::message_callback_wrapper(void** state, struct mqtt_response_publish *published) {
    LwMqtt* mqtt = static_cast<LwMqtt*>(*state);
    if (mqtt) {
        mqtt->handleMessage(published);
    }
}

void LwMqtt::handleMessage(struct mqtt_response_publish *published) {
    if (message_callback_) {
        std::string topic(reinterpret_cast<const char*>(published->topic_name), published->topic_name_size);
        std::string payload(reinterpret_cast<const char*>(published->application_message), published->application_message_size);
        message_callback_(topic, payload);
    }
}

#endif // ENABLE_MQTT